import os
import torch
import numpy as np
from compressai.transforms.functional import yuv_420_to_444, ycbcr2rgb

class RawVideoSequence:
    def __init__(self, file, width, height, seq = None):
        self.file = file
        self.width = width
        self.height = height

        self.sequence = seq

    def __getitem__(self, index):
        raise NotImplementedError()

    def __len__(self) -> int:
        return len(self.sequence)

    def __iter__(self):
        return iter(self.sequence)

    def close(self):
        del self.sequence


class YUV420VideoSequence(RawVideoSequence):
    """yuv420 8bit video sequence"""
    def __init__(self, file, width, height, seq = None):
        super().__init__(file, width, height, seq)


    @staticmethod
    def load_yuv420_video(file, width, height):
        frame_size = int(width * height * 1.5)    
        y_len = width * height
        u_v_len = width * height // 4

        y_planes = []
        u_planes = []
        v_planes = []

        with open(file, "rb") as f:
            while True:
                data = f.read(frame_size)
                if len(data) != frame_size:
                    break

                yuv_data = np.frombuffer(data, dtype = np.uint8)

                y_planes.append(yuv_data[0: y_len]/255.0)
                u_planes.append(yuv_data[y_len: y_len + u_v_len]/255.0)
                v_planes.append(yuv_data[y_len + u_v_len: y_len + u_v_len * 2]/255.0)


        y_planes = torch.Tensor(np.array(y_planes)).view(len(y_planes), 1, height, width)
        u_planes = torch.Tensor(np.array(u_planes)).view(len(u_planes), 1, height // 2, width // 2)
        v_planes = torch.Tensor(np.array(v_planes)).view(len(v_planes), 1, height// 2, width // 2)

        frames = ycbcr2rgb(yuv_420_to_444((y_planes, u_planes, v_planes)))

        return frames
    

    @classmethod
    def from_file(cls, file, width, height) -> "YUV420VideoSequence":
        return cls(
            file = file,
            width = width,
            height = height,
            seq = YUV420VideoSequence.load_yuv420_video(file, width, height)
        )


    @classmethod
    def from_folder(cls, folder, width, height) -> list["YUV420VideoSequence"]:
        sequences = []

        for doc in os.listdir(folder):
            if doc.endswith(".yuv"):
                sequences.append(
                    YUV420VideoSequence(
                        os.path.join(folder, doc),
                        width,
                        height
                    )
                )

        return sequences

    def __getitem__(self, index) -> torch.Tensor:
        if self.sequence is None:
            self.sequence = self.load_yuv420_video(
                self.file, self.width, self.height
            )

        return self.sequence[index]

    def __len__(self) -> int:
        if self.sequence is None:
            self.sequence = self.load_yuv420_video(
                self.file, self.width, self.height
            )
    
        return len(self.sequence)

    def __iter__(self):
        return iter(self.sequence)

    def close(self):
        del self.sequence

